#!/usr/bin/python
# GM_fitc0rho0.py

from time import asctime
from os import linesep
import math
import numpy

import GM_Swap
import GM_apparent
import GM_logN
import GM_N_bvals


# distance # OK?
# iterate # OK
# multiplyarray # OK
# interpolate # OK
# fit # OK
# iterateverbose # OK


def distance(simpars, phypars, compN):
    '''Returns the RMS difference between compN and N_bM.'''
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    # latitudeweightinternal exiled until I know what to think of it
    w = latitudeweightexternal#*latitudeweightinternal
    w2 = w**2
    wn_0 = float(w2.sum())*len(compN)
    dist = N_bM - compN
    wdist = numpy.multiply(dist[:,:], w[numpy.newaxis,:])
    wdist2 = wdist**2
    wdist2_0 = wdist2.sum()
    dist2_0 = wdist2_0/wn_0
    dist_0 = dist2_0**0.5
    return dist_0


def iterate(simpars, phypars, DistMod, ExtMod, distriblist):
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    # Get apparent distribution of stars
    # elements of appearance are 3d (d, l, b)
    appearance = GM_apparent.apparent(simpars, phypars, DistMod, ExtMod, distriblist)
    # Get apparent distribution of stars per degree squared
    # elements of N are 1d (b only)
    # elements of allskyN is 0d, an allsky average
    (N, allskyN) = GM_logN.logN(simpars, phypars, appearance)
    # Condense N to relevant points by linear interpolation or extrapolation
    compN = GM_N_bvals.N_bvals(simpars, phypars, N)
    # And find the 'distance' between the data grid of N values and the simulated grid
    dist = distance(simpars, phypars, compN)

#    del GM_apparent, DistMod, ExtMod, distriblist
#    del plot_logN
#    del GM_N_bvals
    return dist


def multiplyarray(count, factor):
    multarray = numpy.ones(count, dtype=float)
    lmax = 0.5*(count - 1)
    logfactor = factor**(lmax**-1)
#    print 'multarray diagnostics factor lmax logfactor', factor, lmax, logfactor
    for i in range(0, count):
        l = i - 0.5*(count - 1)
        multarray[i] = logfactor**l
    return (logfactor, multarray)


def interpolate(distarray, centralc0, centralrho0, cmultarray, rhomultarray):
    tiny = 1.e-5
    travel = 1.25 # How many half grid lengths does the center shift if local minimum is on edge of grid?
    min = distarray.min()
    print "min is", min
    count = len(cmultarray)
    i0 = 9e9
    j0 = 9e9
    mindiff = 9e9
    carray = centralc0*cmultarray
    rhoarray = centralrho0*rhomultarray
#    cmin = centralc0*cmultarray[0]
#    cmax = centralc0*cmultarray[-1]
#    rhomin = centralrho0*rhomultarray[0]
#    rhomax = centralrho0*rhomultarray[-1]
    for i in range(0, count):
        for j in range(0, count):
            dist = distarray[i,j]
#            print "dist i j is", dist, i, j
            if (dist < min*(1 + tiny)) and (dist > min*(1 - tiny)):
                i0 = int(i)
                j0 = int(j)
            elif (dist - min) < mindiff:
                mindiff = dist - min
    if i0 == 0:
        c0 = centralc0*cmultarray[0]**travel
    elif i0 == (count-1):
        c0 = centralc0*cmultarray[-1]**travel
    elif i0 == 9e9:
        print 'Holy cow, is something wrong with c0 here!'
    else:
        # Parabolic interpolation
        x1 = centralc0*cmultarray[i0-1]
        x2 = centralc0*cmultarray[i0]
        x3 = centralc0*cmultarray[i0+1]
        y1 = distarray[i0-1,j0]
        y2 = distarray[i0,j0]
        y3 = distarray[i0+1,j0]
        a = ((x2-x1)*(y3-y1) + (y2-y1)*(x1-x3))/((x3**2-x1**2)*(x2-x1) + (x3-x1)*(x1**2-x2**2))
        b = (y2 - y1 + a*(x1**2 - x2**2))/(x2-x1)
        c0 = float(-b/(2*a))
#        if (c0 < cmin) or (c0 > cmax):
        if (c0 < carray[i0-1]) or (c0 > carray[i0+1]):
            print 'Your interpolation is wonky! c0 is out of bounds.'
    if j0 == 0:
        rho0 = centralrho0*rhomultarray[0]**travel
    elif j0 == (count-1):
        rho0 = centralrho0*rhomultarray[-1]**travel
    elif j0 == 9e9:
        print 'Holy cow, is something wrong with rho0 here!'
    else:
        # Parabolic interpolation
        x1 = centralrho0*rhomultarray[j0-1]
        x2 = centralrho0*rhomultarray[j0]
        x3 = centralrho0*rhomultarray[j0+1]
        y1 = distarray[i0,j0-1]
        y2 = distarray[i0,j0]
        y3 = distarray[i0,j0+1]
        a = ((x2-x1)*(y3-y1) + (y2-y1)*(x1-x3))/((x3**2-x1**2)*(x2-x1) + (x3-x1)*(x1**2-x2**2))
        b = (y2 - y1 + a*(x1**2 - x2**2))/(x2-x1)
        rho0 = float(-b/(2*a))
#        if (rho0 < rhomin) or (rho0 > rhomax):
        if (rho0 < rhoarray[j0-1]) or (rho0 > rhoarray[j0+1]):
            print 'Your interpolation is wonky! rho0 is out of bounds.'
    return (i0, j0, c0, rho0, mindiff)



#def valmove(fixedval, val0, centralval0, valindex0, count):
#    if fixedval == True:
#        val0 = centralval0
#    elif (valindex == 0) or (valindex == (count - 1)):
#
#                    cfactor = float(clogfactor)**shrink
#                    (clogfactor, cmultarray) = multiplyarray(count, cfactor)
#    return val0

def fit(simpars, phypars, DistMod, ExtMod, distriblist):
    # Input
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    if interpolateMv == True:
        phypars = GM_Swap.swap_fitting(phypars, vanilla=False)
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    # Properties of fit 
    # Condition 1 for break: meet science targets
    targetdiff = 1.
    accuracy = 1.01
    # Condition 2 for break: surpass computation time alotted
    maxiterations = 10
    count = 9 # Should be odd integer five or greater; count**2 is size of 2d array tested
    factor = 1.08 # Total factor searched (mult. up AND down) in 1st iteration
    shrink = 1. # Must be >= 1; larger numbers shrink the array less with successful iteration
    # Obligatory user warning if settings are dumb
    if (fixedrho == True) and (fixedc == True):
        "You used settings that indicated you wanted the density and/or dust extinction constant fit, but also told the program to hold both of those quantities constant.  This is your warning that the program is NOT fitting either density or extinction."
    # Create arrays of values to sample
    (logfactor, multarray) = multiplyarray(count, factor)
    cmultarray = numpy.array(multarray)
    rhomultarray = numpy.array(multarray)
    cfactor = float(factor)
    rhofactor = float(factor)
    clogfactor = float(logfactor)
    rhologfactor = float(logfactor)
    index = 1
    difflist = []
    distlist = []
    distarraylist = []
    clocslist = []
    rlocslist = []
    paramlist = [[c0, rho0, cfactor, rhofactor]]
    # Iterate
    while True:
        print 'Beginning iteration', index, 'over', str(count**2), 'elements. c0 is centered at', c0, 'and rho0 is centered at', rho0, 'with a range of', cfactor, 'and', str(rhofactor)+', respectively, being probed.' 
        print asctime()
        print linesep
        centralc0 = float(c0)
        centralrho0 = float(rho0)
        clocslist.append(centralc0*cmultarray)
        rlocslist.append(centralrho0*rhomultarray)
        cmin = centralc0/cfactor
        cmax = centralc0*cfactor
        rhomin = centralrho0/rhofactor
        rhomax = centralrho0*rhofactor
        distarray = numpy.ones((count,count), dtype=float)
        for i in range(0, count):
            for j in range(0, count):
                simpars['c0'] = cmultarray[i]*centralc0
                simpars['rho0'] = rhomultarray[j]*centralrho0
                dist = iterate(simpars, phypars, DistMod, ExtMod, distriblist)
                distarray[i,j] = float(dist)
#            print 'Line', i, 'of distarray is', distarray[i,:]
        # Test goodness of fit
        (i0, j0, c0, rho0, mindiff) = interpolate(distarray, centralc0, centralrho0, cmultarray, rhomultarray)
        # Record outcome of fit
        ctested = cmultarray[i0]*centralc0
        rhotested = rhomultarray[j0]*centralrho0
        dist = distarray[i0,j0]
        difflist.append(mindiff)
        distlist.append(dist)
        distarraylist.append(distarray)
        # 'Fix' the values of c0, rho0 if needed (no pun intended... ha ha)
        if fixedrho == True:
            rho0 = float(centralrho0)
        if fixedc == True:
            c0 = float(centralc0)
        if (i0 != 0) and (j0 != 0) and (i0 != (count-1)) and (j0 != (count-1)):
            # Decide whether fit is good enough
            if (mindiff < targetdiff) and (clogfactor < accuracy) and (rhologfactor < accuracy):
                print 'Accuracy objectives have been achieved!'
                break
            else:
                # Choose parameters for new iteration
                print 'Neither c0 nor rho0 is on the edge!'
                cfactor = float(clogfactor)**shrink
                (clogfactor, cmultarray) = multiplyarray(count, cfactor)
                rhofactor = float(rhologfactor)**shrink
                (rhologfactor, rhomultarray) = multiplyarray(count, rhofactor)
        elif (i0 != 0) and (i0 != (count-1)) and (fixedrho == True):
            print "The best value for rho0 is on the edge of the sampled region, but since rho0 has a fixed value, the fit in c0 is proceeding nonetheless."
            # Decide whether fit is good enough
            if (mindiff < targetdiff) and (clogfactor < accuracy):
                print 'Accuracy objectives have been achieved!'
                break
            else:
                # Choose parameters for new iteration
                cfactor = float(clogfactor)**shrink
                (clogfactor, cmultarray) = multiplyarray(count, cfactor)
        elif (j0 != 0) and (j0 != (count-1)) and (fixedc == True):
            print "The best value for rho0 is on the edge of the sampled region, but since rho0 has a fixed value, the fit in c0 is proceeding nonetheless."
            # Decide whether fit is good enough
            if (mindiff < targetdiff) and (rhologfactor < accuracy):
                print 'Accuracy objectives have been achieved!'
                break
            else:
                # Choose parameters for new iteration
                rhofactor = float(rhologfactor)**shrink
                (rhologfactor, rhomultarray) = multiplyarray(count, rhofactor)
        # Decide whether fit has gone on too long
        index = index + 1
        if index > maxiterations:
            print 'The maximum number of iterations has been reached.'
            break
        # State for the record
        paramlist.append([c0, rho0, cfactor, rhofactor])
        print linesep
        print 'The best value of dist in this iteration was', dist, 'with indices', i0, j0, 'and parameters c0 =', str(c0)+', rho0 =', str(rho0)+', chosen from ranges', cmin, 'to', str(cmax)+',', rhomin, 'to', str(rhomax)+'. The lowest values within the actual matrix values tested were', ctested, 'and', str(rhotested)+'.'
        print linesep
    paramlist.append([c0, rho0, 0., 0.])
    simpars['c0'] = c0
    simpars['rho0'] = rho0
    print 'I have chosen a value of', rho0, 'for rho0 and a value of', c0, 'for c0, which should be accurate to within a factor of', rhologfactor, 'and', clogfactor, 'respectively.'
    difflist = numpy.array(difflist)
    distlist = numpy.array(distlist)
    distarraylist = numpy.array(distarraylist)
    paramlist = numpy.array(paramlist)
    clocslist = numpy.array(clocslist)
    rlocslist = numpy.array(rlocslist)
    print difflist
    for i in simpars:
        cmd = "simpars['%s'] = %s" % (i,i)
        exec cmd
    return (simpars, difflist, distlist, distarraylist, paramlist, clocslist, rlocslist)



def iterateverbose(simpars, phypars, DistMod, ExtMod, distriblist):
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    # Get apparent distribution of stars
    # elements of appearance are 3d (d, l, b)
    appearance = GM_apparent.apparent(simpars, phypars, DistMod, ExtMod, distriblist)
    # GM_logN and GM_N_bvals produce completely meaningless results in the pecilbeam case, and throw errors on top of it.  So I gate them off here, and make dummy data products below.
    if pencilBeam == False:
        # Get apparent distribution of stars per degree squared
        # elements of N are 1d (b only)
        # elements of allskyN is 0d, an allsky average
        (N, allskyN) = GM_logN.logN(simpars, phypars, appearance)
        # Condense N to relevant points by linear interpolation or extrapolation
        compN = GM_N_bvals.N_bvals(simpars, phypars, N)
    # These dummy products are designed to throw nice errors if unintentionally used.
    else:
        N = 'N is not computed for the pencil beam case.  See fitc0rho0.'
        allskyN = 'allskyN is not computed for the pencil beam case.  See fitc0rho0.'
        compN = 'compN is not computed for the pencil beam case.  See fitc0rho0.'
#    del GM_apparent, DistMod, ExtMod, distriblist
#    del plot_logN
#    del GM_N_bvals
    return (appearance, N, allskyN, compN)
